
import math;



def Dump2DArray(arr):
    dim = len(arr)
    for i in range(dim):
       for j in range(dim):
           print("%5d"%(arr[i][j]),end=' ')
       print(end='\n')


def IsSquare2DArray (arr):
    
    r = len(arr)
    if r == 0:
        return True
    
    
    for i in range(r):
        if len(arr[i]) != r:
            return False
    return True

# @arr: 2D array
# @lst: 1D seq list
# @initDir: 
#         0 indicate LEFTBOTTON
#         1 indicate RIGHTYPPER 
#        [[A B],        
#         [C D]]
#        initDir = 0, output A C B D
#        initDir = 1, output A B C D


def Zigzag(arr,lst, initDir, putToList):
    
    if putToList and IsSquare2DArray(arr) == False:
        return
    
    # list len not x*x number
    list_len = len(lst)
    rt = math.sqrt(list_len)
    if putToList == False and rt.is_integer() == False:
        return
    
    if putToList:
        dim = len(arr)
        lst.clear()
        lst = list(range(dim * dim))
    else:
        arr.clear()
        dim = int(rt)
        print("list len: %d, dim is: %d" %(list_len, dim))
    
        for i in range(dim):
            l = list(range(i * dim, (i + 1) * dim))
            arr.append(l)
        
    
    
    
    r = 0;
    c = 0;
    
    tot = dim * dim
    
    
    goRU =[-1,1]
    goLB =[1,-1]
    
    goDir = []
    goDir.append( goLB )
    goDir.append( goRU )
    LB = 0
    RU = 1
    
    
    dir = initDir
    
    for i in range(tot):
        if putToList:
            list[i] = arr[r][c] 
        else:
            arr[r][c] = lst[i]

        nr = r + goDir[dir][0]
        nc = c + goDir[dir][1] 
        if nr >=0 and nr < dim  and nc >=0 and nc < dim:
            r = nr
            c = nc 
        else:
            
            if dir == LB:
                if r + 1 < dim:
                    r += 1
                else:
                    c += 1
            else:
                if c + 1 < dim:
                    c += 1
                else:
                    r += 1
            dir ^= 1

        
def ZigzagToList(arr, initDir):
    lst=[]
    Zigzag(arr, lst, initDir, True)
    return lst
def ZigzagToArray(lst, initDir):
    arr=[]
    Zigzag(arr, lst, initDir, False)
    return arr

   
   
    
if __name__ == "__main__":


    N = 8
    
    l=list(range(0,N * N))
    l1=list(range(0,4))
    
    arr = ZigzagToArray(l,  1)
    print(arr)
    Dump2DArray(arr)
    
    
    #atest  = ZigzagToArray(l1, 0)
    #btest = ZigzagToArray(l1, 1)
    #Dump2DArray(atest)
    ## output:
    ## 0 2
    ## 1 3
    #Dump2DArray(btest)
    ## output:
    ## 0 1
    ## 2 3